import os.path
import time
import datetime
import pandas as pd
from pymongo import MongoClient
from upload.upload_future import get_main_future_contract_map
from upload.upload_future import get_list_future_code


def load_df_from_mongo(config: dict, contract_dic: dict, ts_code: str):
    sbl = ts_code.split(".")[0][:-4]
    value_li_of_k = contract_dic.setdefault(ts_code.split(".")[0][-2:], [])
    collection_symbol = MongoClient(config['mongodb'])['ts_code_future'][sbl]
    # print("loading ", ts_code)
    df_symbol = pd.DataFrame(collection_symbol.find({"ts_code": ts_code}))
    df_symbol = df_symbol.sort_values(by=['trade_date'], ascending=True)
    # print(ts_code, df_symbol['trade_date'].to_list()[0], df_symbol['trade_date'].to_list()[-1])
    df_symbol.index = df_symbol['trade_date']
    value_li_of_k.append(df_symbol)
    return contract_dic


def load_df_from_mongo_(config: dict, contract_dic: dict, ts_code: str):
    sbl = ts_code.split(".")[0][:-4]
    value_li_of_k = contract_dic.setdefault(ts_code.split(".")[0][-2:], [])
    collection_symbol = MongoClient(config['mongodb'])['tushare_future'][sbl]
    # print("loading ", ts_code)
    print(ts_code)
    get_ = list(collection_symbol.find({"ts_code": ts_code}))
    if len(get_) > 0:
        a = get_[0]
        ts_code = a['ts_code']
        del a['ts_code']
        del a['start_t']
        del a['end_t']
        del a['update']
        df_symbol = pd.DataFrame(a)
        df_symbol = df_symbol.sort_values(by=['t'], ascending=True)
        # print(ts_code, df_symbol['trade_date'].to_list()[0], df_symbol['trade_date'].to_list()[-1])
        df_symbol.index = df_symbol['t']
        df_symbol['ts_code'] = ts_code
        extender = [df_symbol]
    else:
        extender = []
    value_li_of_k.extend(extender)
    return contract_dic


def drop_dup(df):
    t_li, i_li = [], []
    for i, (t, line) in enumerate(df.iterrows()):
        if t not in t_li:
            t_li.append(t)
            i_li.append(i)
    return df.iloc[i_li]


def get_recent_main_codes(code, retro_len):
    recent_main = list(code)[-retro_len:]
    code_dic = {}
    for a in recent_main:
        cd_num = a.split('.')[0][-2:]
        num = code_dic.setdefault(cd_num, [])
        num.append(a)
        print(a)
    unfited_dic = {k: v for k, v in code_dic.items() if len(v) < 3}
    li = list(unfited_dic.values())
    if len(li) > 0:
        for li_ in li:
            re = [s for s in recent_main if s not in li_]
        return re
    else:
        return recent_main


def generate_main_contract_time_series(cfg, ts_code_lower, mode, only_close=True, main_code_li=[]):
    symbol = ts_code_lower.split("_")[0].upper()
    dir_path = "../data/main_contract_cache"
    if os.path.exists(dir_path):
        cache_path = f"{dir_path}/{symbol}.xlsx"
    else:
        cache_path = f"data/main_contract_cache/{symbol}.xlsx"
    if not os.path.exists(cache_path):
        mp = get_main_future_contract_map(cfg)
        list_code = get_list_future_code(cfg)
        listed_code = list_code[symbol]['ts_code']
        code = mp[ts_code_lower].dropna().unique()
        retro_len = max(30, int(0.5 * len(code)))
        recent_main = get_recent_main_codes(code, retro_len)
        # recent_main = list(code)[-retro_len:]
        dic = {}
        re_dic = {}
        for s in recent_main:
            li = dic.setdefault(s.split(".")[0][-2:], [])
            li.append(s)
            re_dic = load_df_from_mongo_(cfg, re_dic, s)
        filter_listed_code = [s for s in listed_code if s.split(".")[0][-2:] in re_dic.keys()]
        for s in filter_listed_code:
            re_dic = load_df_from_mongo_(cfg, re_dic, s)
        df_dic = {k: drop_dup(pd.concat(v)) for k, v in re_dic.items()}
        if only_close:
            df_dic = save_future_excel(df_dic, main_code_li)
        writer = pd.ExcelWriter(cache_path)
        for k, v in df_dic.items():
            v.to_excel(writer, sheet_name=k)
        writer.save()
    else:
        print('load local', cache_path)
        df_dic = {}
        local_dic = pd.read_excel(cache_path, sheet_name=None, index_col=0)
        for k, v in local_dic.items():
            first_col = v.columns.to_list()[0]
            v[first_col] = v[first_col].astype('float')
            df_dic[first_col] = v
    return df_dic, {}


def save_future_excel(dic: dict, main_code_li: list):
    new_dic = {}
    for k_, v_ in dic.items():
        name = v_['ts_code'].to_list()[-1]
        temp_v = v_.loc[:, ['close', 'ts_code']]
        temp_v['close'] = temp_v['close'].astype('float')
        if name.upper() in main_code_li:
            temp_v.columns = [f"{name.upper()}_主力", 'ts_code']
        else:
            temp_v.columns = [f"{name.upper()}", 'ts_code']
        vv = pd.DataFrame(temp_v)
        vv.index = [str(pd.to_datetime(str(s)))[:10] for s in vv.index]
        new_dic[name.upper()] = vv
    return new_dic


if __name__ == "__main__":
    import time
    # 生成期货合约序列
    from tools.utils import load_config, first_second_dir_exists, get_chinese_name_to_ts_code, concat_and_build_dir
    from tools.utils import first_dir_exists

    t1 = time.time()
    path = first_dir_exists()
    cfg = load_config()
    map = get_main_future_contract_map(cfg)
    # for ts_code in map.columns.to_list()[:10]:
    for ts_code in ['ma_zce']:
        data = generate_main_contract_time_series(cfg, ts_code, 'qq')
    t2 = time.time()
    print(t2 - t1)
